#ifndef __M3N_MODEL_H__
#define __M3N_MODEL_H__
/*********************************************************************
 * Software License Agreement (Modified BSD License)
 *
 * Copyright (c) 2009-2010, Willow Garage, Daniel Munoz
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *   * Neither the name of the copyright holders' organizations nor the
 *     names of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *********************************************************************/

#include <omp.h>

#include <algorithm>
#include <string>
#include <iostream>
#include <limits>
#include <vector>
#include <map>
#include <list>
#include <set>

//#include <submodular_graphcut/graphcut_vk3.h>
#include <submodular_graphcut/graphcut_bgl.h> // slower
//#include <submodular_graphcut/graphcut_vk2.h> // very very slow with high order cliques

#include <m3n/random_field.h>
#include <m3n/util/m3n_logging.h>
#include <m3n/regressors/regressor_includes.h>

// --------------------------------------------------------------
//* M3NModel
/*!
 * \brief An abstract class for training Max-Margin Markov Networks
 *        with (Robust) Pott's potentials over arbitrary sized cliques.
 *
 * \see FunctionalM3N
 * \see ParametricM3N
 *
 * This package uses techniques described in the following publications: \n
 * [1] D. Munoz, J. A. Bagnell, N. Vandapel, and M. Hebert,
 *     "Contextual Classification with Functional Max-Margin Markov Networks", CVPR 2009. \n
 * [2] N. Ratliff, D. Silver, and J. A. Bagnell,
 *     "Learning to Search: Functional Gradient Techniques for Imitation Learning", Autonomous Robots 2009 \n
 * [3] B. Taskar, C. Guestrin, and D. Koller, "Max-Margin Markov Networks", NIPS 2003 \n
 * [4] P. Kohli, L. Ladicky, and P. H. S. Torr,
 *     "Robust Higher Order Potentials for Enforcing Label Consistency", IJCV 2009 \n
 * [5] Y. Boykov, O. Veksler, and R. Zabih, "Fast Approximate Energy Minimization via Graph Cuts", PAMI 2001 \n
 * [6] Y. Boykov, and V. Kolmogorov,
 *     "An Experimental Comparison of Min-Cut/Max-Flow Algorithms for Energy Minimization in Vision", PAMI 2004 \n
 * [7] V. Kolmogorov and R. Zabih,  "What Energy Functions can be Minimized via Graph Cuts?", PAMI 2004 \n
 */
// --------------------------------------------------------------
class M3NModel
{
  public:
    // --------------------------------------------------------------
    /*!
     * \brief Instantiates a Max-Margin Markov Network model
     *
     * The model can be used for training with functional gradient boosting,
     * and inference with high-order cliques with robust/associative potentials.
     * The model is forever tied to the specified Robust Pott's truncation parameters
     *
     * \param robust_potts_params Truncation parameters q_i that define the robustness to
     *                            by Q_c = |c|*q_i, for all cliques c in clique set i,
     *                            as defined in Kohli et al. Each value must be in the range [0, 0.5)
     */
    // --------------------------------------------------------------
    M3NModel(const std::vector<double>& robust_potts_params);

    // --------------------------------------------------------------
    /*!
     * \brief Clears this M3NModel so it can be retrained from scratch.
     */
    // --------------------------------------------------------------
    virtual void clear();

    // --------------------------------------------------------------
    /*!
     * \brief Loads this M3N model from file
     *
     * Parameters should be the same from the call to saveToFile
     *
     * \param directory The path containing all model files
     * \param basename The basename of the model files
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int loadFromFile(std::string directory,
                     std::string basename);

    // --------------------------------------------------------------
    /*!
     * \brief Saves this M3N model to file
     *
     * This call will produce a file named <basename>.m3n_model and
     * multiple associated files related to the regressors with the
     * same basename prefix.
     *
     * \param directory The path to store all resulting model files
     * \param basename The basename of the model files
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int saveToFile(std::string directory,
                   std::string basename);

    // --------------------------------------------------------------
    /*!
     * \brief Infers the best labeling for the given RandomField using this trained M3NModel.
     *
     * Edges always follow the Potts model.
     *
     * \param random_field The RandomField to classify, containing Nodes and Cliques
     * \param inferred_labels Mapping node_id -> inferred labels.  If
     * \param max_iterations (Optional) The max number of complete alpha-expansions to perform;
     *                       0 indicates to run until convergence.
     * \param max_nbr_regressors (Optional) The max number of regressors to use;
     *                       0 indicates to use all. (Ignored for ParametricM3N)
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int infer(const RandomField& random_field,
              std::map<unsigned int, unsigned int>& inferred_labels,
              unsigned int max_iterations = 0,
              unsigned int max_nbr_regressors = 0) const;

    // --------------------------------------------------------------
    /*!
     * \brief Trains this M3NModel using the given labeled RandomFields
     *
     * This method can be called multiple times as when doing online learning,
     * but the labels must be present from the previous call
     *
     * \param training_rfs The RandomFields to train upon.
     * \param nbr_iterations Number of iterations to train until
     * \param step_size The learning rate is step_size/sqrt(t)
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int train(const std::vector<const RandomField*>& training_rfs,
              const unsigned int nbr_iterations,
              const double step_size);

    // --------------------------------------------------------------
    /*!
     * \see train
     *
     * This method will save the model every save_interval'th iteration to
     * the specified directory and model basename, as described in saveToFile.
     * The model's basename will be suffixed with "_iterationX" where X is the
     * iteration number it saved at.
     *
     * \param training_rfs The RandomFields to train upon.
     * \param nbr_iterations Number of iterations to train until
     * \param step_size The learning rate is step_size/sqrt(t)
     * \param saving_interval The number of iterations between each saved model
     * \param directory The directory to save the model, see saveToFile
     * \param basename The basename of the the model, see saveToFile
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int train(const std::vector<const RandomField*>& training_rfs,
              const unsigned int nbr_iterations,
              const double step_size,
              unsigned int saving_interval,
              std::string directory,
              std::string basename);

    // --------------------------------------------------------------
    /*!
     * \brief Returns the labels used to train this M3NModel
     */
    // --------------------------------------------------------------
    const std::vector<unsigned int>& getTrainingLabels() const
    {
      return m_training_labels;
    }

  private:
    // --------------------------------------------------------------
    /*!
     * \brief Saves this M3N model to file
     *
     * This method appends basename_suffix to basename for the .m3n_model
     * filename only.  This functionality is used when creating intermediate models
     * during learning
     *
     * \param directory The path to store all resulting model files
     * \param basename The basename of the model files
     * \param basename_suffix An extra suffix appended to the basename of the main
     *                        model file.
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int saveToFile(std::string directory,
                   std::string basename,
                   std::string basename_suffix);

  protected:
    // --------------------------------------------------------------
    /*!
     * \brief Save information specific to the type of trained M3N model
     *
     * \param outfile Stream of base M3NModel file
     * \param directory The path to store all resulting model files
     * \param basename The basename of the model files
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    virtual int doSaveToFile(std::ofstream& outfile,
                             std::string directory,
                             std::string basename) = 0;

    // --------------------------------------------------------------
    /*!
     * \brief Loads information specific to the type of trained M3N model
     *
     * \param infile Stream of base M3NModel file
     * \param directory The path to store all resulting model files
     * \param basename The basename of the model files
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    virtual int doLoadFromFile(std::ifstream& infile,
                               std::string directory,
                               std::string basename) = 0;

    // --------------------------------------------------------------
    /*!
     * \brief Initialize stacked feature indices information
     */
    // --------------------------------------------------------------
    void initStackedFeatureIndices();

    // ===================================================================
    /*! \name Learning related  */
    // ===================================================================
    //@{
    // --------------------------------------------------------------
    /*!
     * \brief Performs a subgradient update
     *
     * \param training_rfs The list of RandomFields to train on
     * \param infer_random  Flag if to use random inference (used on first iteration)
     * \param step_size How much to move in subgradient direction
     * \param logger Container to hold timing information
     */
    // --------------------------------------------------------------
    virtual void doSubgradientUpdate(const std::vector<const RandomField*>& training_rfs,
                                     const bool infer_random,
                                     const double step_size,
                                     M3NLogger& logger) = 0;

    // --------------------------------------------------------------
    /*!
     * \brief Computes the functional gradient residuals of the structured loss function
     *
     * \param training_rfs The list of RandomFields to train on
     * \param infer_random  Flag if to use random inference (used on first iteration)
     * \param regressor Regressor to hold residual target values
     * \param logger Container to hold timing information
     */
    // --------------------------------------------------------------
    void structLossResiduals(const std::vector<const RandomField*>& training_rfs,
                             const bool infer_random,
                             RegressorWrapper& regressor,
                             M3NLogger& logger);

    // --------------------------------------------------------------
    /*!
     * \brief Verifies the given RandomFields have consistent labels and
     *        feature dimensions for training
     *
     * \param training_rfs The list of RandomFields
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int extractVerifyLabelsFeatures(const std::vector<const RandomField*>& training_rfs);

    // --------------------------------------------------------------
    /*!
     * \brief Computes the functional gradient residual for high-order cliques.
     *
     * \param truncation_param The Robust Potts truncation factor for the clique
     * \param clique_order The order of the clique
     * \param nbr_mode_label The number of nodes in the clique containing the mode label
     *
     * \return The functional gradient residual
     */
    // --------------------------------------------------------------
    double calcFuncGradResidual(const double truncation_param,
                                const unsigned int clique_order,
                                const unsigned int nbr_mode_label) const;
    //@}

    // ===================================================================
    /*! \name Inference related  */
    // ===================================================================
    //@{
    // --------------------------------------------------------------
    /*!
     * \brief Inference without ensuring the model is already trained
     *
     * \see infer
     */
    // --------------------------------------------------------------
    int doInfer(const RandomField& random_field,
                std::map<unsigned int, unsigned int>& inferred_labels,
                unsigned int max_iterations = 0,
                unsigned int max_nbr_regressors = 0) const;

    // --------------------------------------------------------------
    /*!
     * \brief Computes and saves all the potential values for assigning all
     *        possible training labels to each Node and Clique in the given
     *        RandomField
     *
     * \param random_field The RandomField to compute all potentials for
     * \param cache_node_potentials node id -> label -> value
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int
        cachePotentials(const RandomField& random_field,
                        const unsigned int max_nbr_regressors,
                        std::map<unsigned int, std::map<unsigned int, double> >& cache_node_potentials,
                        std::vector<std::map<unsigned int, std::map<unsigned int, double> > >& cache_clique_set_potentials) const;

    // --------------------------------------------------------------
    /*!
     * \brief Generates an arbitrary initial labeling for alpha-expansion
     *
     * Either random or best node labels (see implementation)
     *
     * \param random_field The RandomField that will be performing inference on
     * \param inferred_labels Mapping of node ids to initial label value
     */
    // --------------------------------------------------------------
    void
        generateInitialLabeling(const RandomField& random_field,
                                std::map<unsigned int, unsigned int>& inferred_labels,
                                const std::map<unsigned int, std::map<unsigned int, double> >& cache_node_potentials) const;

    // --------------------------------------------------------------
    /*!
     * \brief Computes the potential value for assigning the node the specified label
     *
     * \param node The node containing its features
     * \param label The label to assign node
     * \param potential_val Reference to hold the computed potential value
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    virtual int computePotential(const RandomField::Node& node,
                                 const unsigned int label,
                                 const size_t max_nbr_regressors,
                                 double& potential_val) const = 0;

    // --------------------------------------------------------------
    /*!
     * \brief Computes the potential value for assigning all the nodes within the
     *        clique the specified label
     *
     * \param clique The clique containing its features
     * \param clique_set_idx The clique set index of clique
     * \param label The label to assign all the nodes in the clique
     * \param potential_val Reference to hold the computed potential value
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    virtual int computePotential(const RandomField::Clique& clique,
                                 const unsigned int clique_set_idx,
                                 const unsigned int label,
                                 const size_t max_nbr_regressors,
                                 double& potential_val) const = 0;

    // --------------------------------------------------------------
    /*!
     * \brief Adds node energy term for alpha-expansion graph-cut inference
     *
     * \param node The node containing its features
     * \param energy_func The submodular energy function minimizer
     * \param energy_var The variable in the energy minimizer that represents node
     * \param curr_label The current node's current label
     * \param alpha_label The alpha label to potentially expand to
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int
        addNodeEnergy(const RandomField::Node& node,
                      submodular_graphcut::SubmodularGraphcut& energy_func,
                      const submodular_graphcut::SubmodularGraphcut::EnergyVar& energy_var,
                      const unsigned int curr_label,
                      const unsigned int alpha_label,
                      const std::map<unsigned int, std::map<unsigned int, double> >& cache_node_potentials) const;

    // --------------------------------------------------------------
    /*!
     * \brief Adds edge (pairwise) energy term for alpha-expansion graph-cut inference
     *
     * \param edge The edge containing its features
     * \param clique_set_idx The clique set that the edge belongs to
     * \param energy_func The submodular energy function minimizer
     * \param energy_vars All variables in the energy minimizer that represents the RandomField
     * \param curr_labeling The RandomField's current labeling
     * \param alpha_label The alpha label to potentially expand to
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int
        addEdgeEnergy(const RandomField::Clique& edge,
                      submodular_graphcut::SubmodularGraphcut& energy_func,
                      const std::map<unsigned int,
                          submodular_graphcut::SubmodularGraphcut::EnergyVar>& energy_vars,
                      const std::map<unsigned int, unsigned int>& curr_labeling,
                      const unsigned int alpha_label,
                      const std::map<unsigned int, std::map<unsigned int, double> >& cache_cs_potentials) const;

    // --------------------------------------------------------------
    /*!
     * \brief Adds Pn Potts potential term for alpha-expansion graph-cut inference
     *
     * \param clique The high-order clique containing its features
     * \param clique_set_idx The clique set that the edge belongs to
     * \param energy_func The submodular energy function minimizer
     * \param energy_vars All variables in the energy minimizer that represents the RandomField
     * \param curr_labeling The RandomField's current labeling
     * \param alpha_label The alpha label to potentially expand to
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int
        addCliqueEnergyPotts(const RandomField::Clique& clique,
                             submodular_graphcut::SubmodularGraphcut& energy_func,
                             const std::map<unsigned int,
                                 submodular_graphcut::SubmodularGraphcut::EnergyVar>& energy_vars,
                             const std::map<unsigned int, unsigned int>& curr_labeling,
                             const unsigned int alpha_label,
                             const std::map<unsigned int, std::map<unsigned int, double> >& cache_cs_potentials) const;

    // --------------------------------------------------------------
    /*!
     * \brief Adds Robust Pn Potts potential term for alpha-expansion graph-cut inference
     *
     * \param clique The high-order clique containing its features
     * \param clique_set_idx The clique set that the edge belongs to
     * \param energy_func The submodular energy function minimizer
     * \param energy_vars All variables in the energy minimizer that represents the RandomField
     * \param curr_labeling The RandomField's current labeling
     * \param alpha_label The alpha label to potentially expand to
     *
     * \return 0 on success, otherwise negative value on error
     */
    // --------------------------------------------------------------
    int
        addCliqueEnergyRobustPotts(const RandomField::Clique& clique,
                                   const unsigned int clique_set_idx,
                                   submodular_graphcut::SubmodularGraphcut& energy_func,
                                   const std::map<unsigned int,
                                       submodular_graphcut::SubmodularGraphcut::EnergyVar>& energy_vars,
                                   const std::map<unsigned int, unsigned int>& curr_labeling,
                                   const unsigned int alpha_label,
                                   const std::map<unsigned int, std::map<unsigned int, double> >& cache_cs_potentials) const;
    //@}

    bool m_trained;
    bool m_loss_augmented_inference;

    size_t m_node_feature_dim;
    std::vector<size_t> m_clique_set_feature_dims;

    std::map<unsigned int, size_t> m_node_stacked_feature_start_idx;
    std::vector<std::map<unsigned int, size_t> > m_clique_set_stacked_feature_start_idx;
    size_t m_total_stack_feature_dim;

    std::vector<double> m_robust_potts_params;
    std::vector<unsigned int> m_training_labels;

    // Number of threads to use for caching clique potentials
    // This should only be 1 during training
    int m_nbr_threads_cache;
};

#endif
